import logging
import os
import json
import time
import traceback
from typing import Dict, List, Optional, Any

# 导入OpenAI客户端库
try:
    import openai
    OPENAI_AVAILABLE = True
except ImportError:
    OPENAI_AVAILABLE = False

# 配置日志
logger = logging.getLogger(__name__)

# 配置管理
class RepairSuggestionConfig:
    """修复建议服务配置"""
    def __init__(self):
        # 从环境变量获取API密钥，优先使用OpenAI
        self.openai_api_key = os.environ.get('OPENAI_API_KEY')
        
        # 默认配置
        self.model = "gpt-4o-mini"  # 使用较经济的模型
        self.max_tokens = 1000
        self.temperature = 0.2  # 较低温度保证输出的一致性和确定性
        self.top_p = 0.9
        self.frequency_penalty = 0.0
        self.presence_penalty = 0.0
        
        # 支持的修复建议类型
        self.supported_vuln_types = [
            'sql_injection', 'xss', 'command_injection', 'file_upload',
            'path_traversal', 'csrf', 'ssrf', 'info_leak',
            'authentication_bypass', 'authorization_bypass', 'ssti',
            'insecure_deserialization', 'xxe', 'weak_password_policy',
            'security_headers', 'insecure_cookies', 'missing_authentication',
            'dependency_vulnerability'
        ]
        
        # 缓存配置
        self.cache_enabled = True
        self.cache_dir = os.path.join(os.path.dirname(__file__), '..', 'cache')
        self.cache_ttl = 3600 * 24 * 7  # 缓存有效期7天
        
        # 重试配置
        self.max_retries = 3
        self.retry_delay = 2
        
        # 初始化缓存目录
        if self.cache_enabled and not os.path.exists(self.cache_dir):
            os.makedirs(self.cache_dir)

# 缓存管理
class SuggestionCache:
    """修复建议缓存管理器"""
    def __init__(self, config: RepairSuggestionConfig):
        self.config = config
        self.cache_dir = config.cache_dir
    
    def get_cache_key(self, vuln_info: Dict[str, Any]) -> str:
        """生成缓存键"""
        # 使用漏洞类型、路径、参数等关键信息生成缓存键
        key_parts = [
            vuln_info.get('vuln_type', 'unknown'),
            vuln_info.get('path', '').replace('/', '_'),
            vuln_info.get('parameter', '').replace('=', '_'),
            str(vuln_info.get('severity', 'medium'))
        ]
        return '_'.join([part[:30] for part in key_parts if part])[:100]
    
    def get(self, vuln_info: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """从缓存获取建议"""
        if not self.config.cache_enabled:
            return None
        
        try:
            cache_key = self.get_cache_key(vuln_info)
            cache_file = os.path.join(self.cache_dir, f"suggestion_{cache_key}.json")
            
            if not os.path.exists(cache_file):
                return None
            
            # 检查缓存是否过期
            if time.time() - os.path.getmtime(cache_file) > self.config.cache_ttl:
                os.remove(cache_file)  # 删除过期缓存
                return None
            
            with open(cache_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            logger.warning(f"Cache read error: {str(e)}")
            return None
    
    def set(self, vuln_info: Dict[str, Any], suggestion: Dict[str, Any]):
        """保存建议到缓存"""
        if not self.config.cache_enabled:
            return
        
        try:
            cache_key = self.get_cache_key(vuln_info)
            cache_file = os.path.join(self.cache_dir, f"suggestion_{cache_key}.json")
            
            # 添加缓存元数据
            suggestion['cached_at'] = time.time()
            
            with open(cache_file, 'w', encoding='utf-8') as f:
                json.dump(suggestion, f, ensure_ascii=False, indent=2)
        except Exception as e:
            logger.warning(f"Cache write error: {str(e)}")

# 修复建议服务
class RepairSuggestionService:
    """漏洞修复建议服务"""
    def __init__(self):
        self.config = RepairSuggestionConfig()
        self.cache = SuggestionCache(self.config)
        
        # 初始化OpenAI客户端
        if OPENAI_AVAILABLE and self.config.openai_api_key:
            openai.api_key = self.config.openai_api_key
            logger.info("OpenAI client initialized successfully")
        else:
            logger.warning("OpenAI client not available, will use fallback suggestions")
        
        # 预定义的修复建议模板
        self.fallback_suggestions = {
            'sql_injection': {
                'title': 'SQL注入漏洞修复建议',
                'description': 'SQL注入是一种常见的高危漏洞，攻击者可以通过注入恶意SQL代码来操作数据库。',
                'fix_steps': [
                    '使用参数化查询（Prepared Statements）替代字符串拼接',
                    '实施输入验证和净化',
                    '使用ORM框架进行数据库操作',
                    '实施最小权限原则',
                    '添加WAF规则防护'
                ],
                'code_example': """# 不安全的实现
query = "SELECT * FROM users WHERE username = '" + username + "'";

# 安全的实现 - 参数化查询
query = "SELECT * FROM users WHERE username = ?";
statement = connection.prepareStatement(query);
statement.setString(1, username);""",
                'tools': ['SQLMap', 'OWASP ZAP', 'ESLint（SQL规则）']
            },
            'xss': {
                'title': '跨站脚本（XSS）漏洞修复建议',
                'description': 'XSS漏洞允许攻击者在用户浏览器中执行恶意脚本，可能导致会话劫持、数据窃取等。',
                'fix_steps': [
                    '对所有用户输入进行严格验证和过滤',
                    '对输出到HTML的内容进行HTML实体编码',
                    '实施内容安全策略（CSP）',
                    '使用现代框架的自动转义功能',
                    '验证HTTP Only和Secure Cookie设置'
                ],
                'code_example': """// 不安全的实现
app.get('/user', (req, res) => {
  res.send('Hello ' + req.query.name);
});

// 安全的实现
const escapeHtml = (unsafe) => {
  return unsafe
    .replace(/&/g, "&amp;")
    .replace(/</g, "&lt;")
    .replace(/>/g, "&gt;")
    .replace(/"/g, "&quot;")
    .replace(/'/g, "&#039;");
};

app.get('/user', (req, res) => {
  res.send('Hello ' + escapeHtml(req.query.name));
});""",
                'tools': ['DOMPurify', 'CSP Builder', 'XSS Hunter']
            },
            'command_injection': {
                'title': '命令注入漏洞修复建议',
                'description': '命令注入允许攻击者在目标系统上执行恶意命令，可能导致系统完全被接管。',
                'fix_steps': [
                    '使用语言内置的安全函数替代直接的系统命令调用',
                    '实施严格的输入验证和白名单机制',
                    '使用参数数组方式执行命令而非字符串拼接',
                    '限制执行命令的进程权限',
                    '避免在web应用中执行系统命令'
                ],
                'code_example': """# 不安全的实现
import os
os.system(f'ping -c 4 {target}')

# 安全的实现
import subprocess
try:
    subprocess.run(['ping', '-c', '4', target], check=True, capture_output=True, text=True)
except subprocess.CalledProcessError as e:
    print(f"Command failed: {e}")""",
                'tools': ['OSSEC', 'Auditd', 'ModSecurity']
            }
        }
    
    def generate_prompt(self, vulnerability: Dict[str, Any]) -> str:
        """生成高质量的提示词"""
        vuln_type = vulnerability.get('vuln_type', 'unknown')
        path = vulnerability.get('path', '')
        parameter = vulnerability.get('parameter', '')
        payload = vulnerability.get('payload', '')
        evidence = vulnerability.get('evidence', '')
        severity = vulnerability.get('severity', 'medium')
        details = vulnerability.get('details', {})
        
        # 技术栈信息
        tech_stack = vulnerability.get('tech_stack', [])
        if not tech_stack:
            # 尝试从路径推断技术栈
            if any(ext in path for ext in ['.php', '.phtml']):
                tech_stack.append('PHP')
            elif any(ext in path for ext in ['.js', '.ts', '.jsx', '.tsx']):
                tech_stack.append('JavaScript')
                tech_stack.append('Node.js')
            elif any(ext in path for ext in ['.py']):
                tech_stack.append('Python')
                if 'django' in path.lower() or '/admin/' in path:
                    tech_stack.append('Django')
                elif 'flask' in path.lower():
                    tech_stack.append('Flask')
            elif any(ext in path for ext in ['.java']):
                tech_stack.append('Java')
        
        tech_stack_str = ', '.join(tech_stack) if tech_stack else '未知'
        
        # 构建详细的提示词
        prompt = f"""你是一位专业的网络安全工程师，请为以下漏洞提供详细、准确且可操作的修复建议。

## 漏洞信息
- 漏洞类型: {vuln_type}
- 严重程度: {severity}
- 发现路径: {path}
- 影响参数: {parameter}
- 测试载荷: {payload}
- 发现证据: {evidence}
- 技术栈: {tech_stack_str}

## 请提供的修复建议内容
1. 漏洞详细解释（简短清晰，不超过50字）
2. 修复步骤（3-5点，简洁明了，可直接操作）
3. 代码示例（使用{tech_stack_str}相关语言，展示不安全的实现和安全的实现对比）
4. 相关安全最佳实践建议
5. 推荐使用的工具或库

## 回答格式要求
请按照以下JSON格式返回，确保JSON可以正确解析，不要包含任何格式说明文字：
{"title": "修复建议标题", "description": "漏洞简短解释", "fix_steps": ["步骤1", "步骤2", "步骤3", "步骤4", "步骤5"], "code_example": "安全vs不安全代码示例", "best_practices": ["最佳实践1", "最佳实践2", "最佳实践3"], "tools": ["工具1", "工具2", "工具3"]}

请确保回答直接针对该漏洞，代码示例与提供的技术栈相关，建议具体且可操作。"""
        
        return prompt
    
    def get_ai_suggestion(self, vulnerability: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """使用AI获取修复建议"""
        if not OPENAI_AVAILABLE or not self.config.openai_api_key:
            logger.warning("OpenAI not available, skipping AI suggestion")
            return None
        
        prompt = self.generate_prompt(vulnerability)
        
        for attempt in range(self.config.max_retries):
            try:
                # 调用OpenAI API
                response = openai.chat.completions.create(
                    model=self.config.model,
                    messages=[
                        {"role": "system", "content": "你是一位专业的网络安全工程师，擅长提供详细且可操作的漏洞修复建议。"},
                        {"role": "user", "content": prompt}
                    ],
                    max_tokens=self.config.max_tokens,
                    temperature=self.config.temperature,
                    top_p=self.config.top_p,
                    frequency_penalty=self.config.frequency_penalty,
                    presence_penalty=self.config.presence_penalty
                )
                
                # 解析响应
                content = response.choices[0].message.content.strip()
                
                # 尝试解析JSON
                try:
                    suggestion = json.loads(content)
                    logger.info(f"Successfully generated AI suggestion for {vulnerability.get('vuln_type')}")
                    return suggestion
                except json.JSONDecodeError:
                    logger.warning(f"Failed to parse AI response as JSON, using fallback")
                    # 尝试提取JSON部分
                    if '{' in content and '}' in content:
                        json_part = content[content.find('{'):content.rfind('}')+1]
                        try:
                            return json.loads(json_part)
                        except:
                            pass
                    return None
                    
            except Exception as e:
                logger.error(f"OpenAI API error (attempt {attempt+1}/{self.config.max_retries}): {str(e)}")
                if attempt < self.config.max_retries - 1:
                    time.sleep(self.config.retry_delay * (attempt + 1))  # 指数退避
                else:
                    logger.error(f"All retries failed for OpenAI API: {traceback.format_exc()}")
        
        return None
    
    def get_fallback_suggestion(self, vulnerability: Dict[str, Any]) -> Dict[str, Any]:
        """获取回退修复建议"""
        vuln_type = vulnerability.get('vuln_type', 'unknown')
        
        # 获取预定义的修复建议
        if vuln_type in self.fallback_suggestions:
            suggestion = self.fallback_suggestions[vuln_type].copy()
        else:
            # 通用修复建议
            suggestion = {
                'title': f'{vuln_type}漏洞修复建议',
                'description': '该漏洞需要进行专业安全修复，请遵循安全最佳实践。',
                'fix_steps': [
                    '进行详细的安全评估',
                    '根据漏洞类型实施相应的修复措施',
                    '实施输入验证和输出编码',
                    '遵循最小权限原则',
                    '进行修复后的安全测试'
                ],
                'code_example': '# 请根据具体漏洞类型和技术栈实施相应的安全修复\n# 一般原则：验证所有输入，编码所有输出，使用安全的API',
                'tools': ['OWASP ZAP', 'Burp Suite', 'Nessus']
            }
        
        # 添加最佳实践字段
        if 'best_practices' not in suggestion:
            suggestion['best_practices'] = [
                '定期进行安全代码审查',
                '实施安全开发生命周期(SDLC)',
                '保持软件和依赖组件的更新',
                '进行定期的安全培训',
                '使用自动化安全测试工具'
            ]
        
        return suggestion
    
    def get_repair_suggestion(self, vulnerability: Dict[str, Any]) -> Dict[str, Any]:
        """获取单个漏洞的修复建议"""
        try:
            # 检查缓存
            cached_suggestion = self.cache.get(vulnerability)
            if cached_suggestion:
                logger.info(f"Using cached suggestion for {vulnerability.get('vuln_type')}")
                return cached_suggestion
            
            vuln_type = vulnerability.get('vuln_type', 'unknown')
            
            # 检查是否支持该漏洞类型
            if vuln_type not in self.config.supported_vuln_types:
                logger.warning(f"Vulnerability type {vuln_type} not supported")
                suggestion = self.get_fallback_suggestion(vulnerability)
            else:
                # 优先使用AI建议
                suggestion = self.get_ai_suggestion(vulnerability)
                
                # 如果AI建议失败，使用回退建议
                if not suggestion:
                    logger.info(f"Using fallback suggestion for {vuln_type}")
                    suggestion = self.get_fallback_suggestion(vulnerability)
            
            # 添加元数据
            suggestion['source'] = 'ai' if suggestion.get('source') == 'ai' else 'fallback'
            suggestion['timestamp'] = time.time()
            suggestion['vulnerability_id'] = vulnerability.get('id', 'unknown')
            
            # 缓存结果
            self.cache.set(vulnerability, suggestion)
            
            return suggestion
            
        except Exception as e:
            logger.error(f"Error generating repair suggestion: {str(e)}")
            # 返回最小化的错误建议
            return {
                'title': '修复建议生成失败',
                'description': f'无法生成详细修复建议: {str(e)}',
                'fix_steps': ['请联系安全团队获取详细修复指导'],
                'source': 'error',
                'error': True
            }
    
    def batch_get_repair_suggestions(self, vulnerabilities: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """批量获取修复建议"""
        results = []
        
        try:
            logger.info(f"Processing batch repair suggestions for {len(vulnerabilities)} vulnerabilities")
            
            # 批量处理，优化API调用和缓存使用
            for vulnerability in vulnerabilities:
                suggestion = self.get_repair_suggestion(vulnerability)
                results.append({
                    'vulnerability_id': vulnerability.get('id', f'vuln_{len(results)}'),
                    'suggestion': suggestion
                })
                
            logger.info(f"Batch repair suggestions completed: {len(results)} suggestions generated")
            
        except Exception as e:
            logger.error(f"Error in batch repair suggestions: {str(e)}")
            # 返回错误信息
            for vulnerability in vulnerabilities:
                results.append({
                    'vulnerability_id': vulnerability.get('id', f'vuln_{len(results)}'),
                    'suggestion': {
                        'title': '批量处理失败',
                        'description': f'批量处理时出错: {str(e)}',
                        'error': True
                    }
                })
        
        return results

# 全局服务实例
repair_suggestion_service = RepairSuggestionService()

# 便捷函数
def get_repair_suggestion(vulnerability: Dict[str, Any]) -> Dict[str, Any]:
    """获取单个漏洞的修复建议（便捷函数）"""
    return repair_suggestion_service.get_repair_suggestion(vulnerability)

def batch_get_repair_suggestions(vulnerabilities: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """批量获取修复建议（便捷函数）"""
    return repair_suggestion_service.batch_get_repair_suggestions(vulnerabilities)